/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"
#include <iostream>

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
    treeHeight = 0;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	// Try to open the B+ tree page file
	RC result = pf.open(indexname, mode);
	if( result != 0)
	{
		return result;
	}
	
	// Save the mode that we opened the file in
	index_mode = mode;
	
	// If in read mode, obtain the root pid and the height of the tree
	if(index_mode == 'r')
	{
		char buf[PageFile::PAGE_SIZE];
		char* iter = &buf[0];
		result = pf.read(0, buf);
		if(result != 0)
		{
			return result;
		}
		
		memcpy(&rootPid, iter, sizeof(PageId));
		iter += sizeof(PageId);
		memcpy(&treeHeight, iter, sizeof(int));		
	}
	
    return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
	// If in write mode, save the root PID and the tree height 
	// into the index file
	
	if(index_mode == 'w')
	{
		char buf[PageFile::PAGE_SIZE];
		char* iter = &buf[0];
		memcpy(iter, &rootPid, sizeof(PageId));
		iter += sizeof(PageId);
		memcpy(iter, &treeHeight, sizeof(int));
		if(pf.write(0, buf) != 0)
		{
			return RC_FILE_WRITE_FAILED;
		}
	}
	
	// Close the B+ tree page file
	pf.close();
	
    return 0;
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
  RC rc;
	
	// No trees, currently no data at all, create a leaf for root
	if(rootPid == -1)	{
    // Put the first pair into the root
		BTLeafNode firstRoot;
    if ((rc = firstRoot.insert(key, rid)) != 0) {
      return rc;
    }
    // Get a new pid for the new root, and save it
    rootPid = pf.endPid();
    // Write to the pf, make the pf allocate that page to the root
    firstRoot.write(rootPid, pf);
    // The tree now has a valid height
    treeHeight = 1;
    return 0;
	}
	// If the tree has more than one level, seek the correct leaf node
	else if(treeHeight > 0)	{
    int keyFromBelow;
    PageId pageFromBelow;
    // Recursively insert into the tree
		rc = insertBelow(rootPid, treeHeight, key, rid, keyFromBelow, pageFromBelow);
    
    // The root node overflowed, create a new one
    if(rc == RC_NODE_FULL) {
      // Create an empty node for the new root
        BTNonLeafNode newRoot;
        // The old root on the left, and its new sibling on the right
        newRoot.initializeRoot(rootPid, keyFromBelow, pageFromBelow);
        // Get a new page from the pagefile, and write the new root to it
        PageId newPid = pf.endPid();
        if((rc = newRoot.write(pf.endPid(), pf)) != 0) {
          // Write error
          return rc;
        }
        // Save the fact that there's new root, and increased height of the tree
        treeHeight++;
        rootPid = newPid;
        return 0;
    }
    // The value was inserted without overflow, or error
    else {
      return rc;
    }
	}
  return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	// Create a page and entry ID tracker
	PageId read_pid = rootPid;
	int entry_id = 0;
	
	// Check for empty tree
	if(treeHeight <= 0)
	{
		// Not sure what to do...
	}
	
	// If the tree has more than one level, seek the correct leaf node
	else if(treeHeight > 1)
	{
		// Initialize a Non-Leaf node and pid tracker
		BTNonLeafNode nl_node;
	
		// Make our way to the leaf level of our tree
		for(int k = 1; k < treeHeight; k++)
		{
			// Load the node data, and check for errors
			if(nl_node.read(read_pid, pf) != 0)
			{
				return RC_FILE_READ_FAILED;
			}
			
			// Look for the correct child, and check for errors
			if(nl_node.locateChildPtr(searchKey, read_pid) != 0)
			{
				return RC_INVALID_PID;
			}
		}
	}
	
	// Initialize the leaf node
	BTLeafNode l_node;
	
	// Load the node data, and check for errors
	if(l_node.read(read_pid, pf) != 0)
	{
		return RC_FILE_READ_FAILED;
	}
	
	// Get the pid of the value, and check for errors
	if(l_node.locate(searchKey, entry_id))
	{
		return RC_INVALID_PID;
	}
	
	// Save values into the cursor
	cursor.pid = read_pid;
	cursor.eid = entry_id;
	
    return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in 
						 the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	// Create variables that will cursor through the trees
	BTLeafNode l_node;
	PageId pid_reader = cursor.pid;
	int	eid_reader = cursor.eid;
	
	// Read in the first leaf, and check for success
	if(l_node.read(pid_reader, pf) != 0)
	{
		return RC_FILE_READ_FAILED;
	}
	
	// Read the first entry
	if(l_node.readEntry(eid_reader, key, rid) != 0)
	{
		return RC_INVALID_RID;
	}
	
	// Increment eid, and check for end of leaf
	// If you finish reading the node, reset eid and increment pid
	eid_reader++;
	if(eid_reader >= l_node.getKeyCount())
	{
		eid_reader = 0;
		pid_reader = l_node.getNextNodePtr();
		
		// If this is the last leaf in the tree, return an end of tree error
		if(pid_reader == -1)
		{
			return RC_END_OF_TREE;
		}
	}
	
	// Increment the cursor, and return
	cursor.pid = pid_reader;
	cursor.eid = eid_reader;
	
    return 0;
}

  /*
  * Insert the key below recursively, get back a midkey and page key if needed
  * and needs to be inserted into the current level.  If current level overflows
  * Then return the pid of the new node and the midKey that is passed to the parent.
  * @param currNode[IN] the pid of the node being considered for insertion.
  * @param height[IN] the height of the subtree that includes currNode and below.
  * @param key[IN] the key for the value inserted into the index
  * @param rid[IN] the RecordId for the record being inserted into the index
  * @param midKey[OUT] the middle key from below to be inserted to the parent
  * @param pid[OUT] the pid of the new node to be inserted into the parent
  * @return error code. 0 if no error.  Returns node full if has values to pass up
  */
RC BTreeIndex::insertBelow(PageId currNodePage, int height, int key, const RecordId& rid, int& midKey, PageId& pid)
{
  RC rc;
  PageId nextNodePage;
  int keyFromBelow;
  PageId pidFromBelow;
  // Somewhere in the middle of the tree
  if(height > 1) {
    BTNonLeafNode currNode;
    // Read in the current node to examine
    if((rc = currNode.read(currNodePage, pf)) != 0)	{
			return rc;
    }
    
    // Look for the correct child, and check for errors
    if((rc = currNode.locateChildPtr(key, nextNodePage)) != 0) {
      return rc;
		}
    // Attempt to have the key inserted into that subtree, recursion!
    rc = insertBelow(nextNodePage, height - 1, key, rid, keyFromBelow, pidFromBelow);
    
    // Success, no overflow, simply return
    if(rc == 0) {
      return 0;
    }
    // The node was full, so some overflow occured
    else if(rc == RC_NODE_FULL) {
      // Attempt to insert the pointer to the new node
      rc = currNode.insert(keyFromBelow, pidFromBelow);
      
      if(rc == 0) {
        rc = currNode.write(currNodePage, pf);
        return rc;
      }
      // If this node is full as well, insert, split and propagate up
      else if(rc == RC_NODE_FULL) {
        // Create an empty node for the new sibling
        BTNonLeafNode sibling;
        // Insert and split, getting back midKey to insert in parent
        currNode.insertAndSplit(keyFromBelow, pidFromBelow, sibling, midKey);
        // Save the new sibling node into the pf, a new page for a new node
        pid = pf.endPid();
        if((rc = sibling.write(pf.endPid(), pf)) != 0) {
          // Write error
          return rc;
        }
        
        
        if((rc = currNode.write(currNodePage, pf)) != 0)
        {
          return rc;
        }
        return RC_NODE_FULL;
      }
      // This node was not full after value was inserted
      else {
        return rc;
      }
    }
    // An error occcured, propagate that error
    else
      return rc;
  }
  // Otherwise, at a leaf node, final recursion step before returning
  else if(height == 1) {
    BTLeafNode l_node;
    // Load the node data, and check for errors
    if((rc = l_node.read(currNodePage, pf)) != 0)
    {
      return rc;
    }
    // Attempt to insert the value
    rc = l_node.insert(key, rid);
    // If the leaf is full, split and propagate up
    if(rc == RC_NODE_FULL){
      // Create an empty node for the new sibling
      BTLeafNode sibling;
      // Insert and split, getting back key to pass up
      rc = l_node.insertAndSplit(key, rid, sibling, midKey);
      // Save the new sibling node into the pf, a new page for a new node
      pid = pf.endPid();
      
      // Set the next ptrs for both leaf nodes
      sibling.setNextNodePtr(l_node.getNextNodePtr());
      l_node.setNextNodePtr(pid);
      
      if((rc = sibling.write(pf.endPid(), pf)) != 0) {
        // Write error
        return rc;
      }
      
      // Also write the other node!
      if((rc = l_node.write(currNodePage, pf)) != 0)
      {
        return rc;
      }
      // Since we overflowed, pass the value up
      return RC_NODE_FULL;
    }
    // The leaf was not full when value was inserted
    else {
      if(rc < 0)
        return rc;
      rc = l_node.write(currNodePage, pf);
      return rc;
    }
  }
  // Otherwise some kind of invalid height, we reached negative or zero height
  return RC_END_OF_TREE;
}
